''' Helper functions that are use by, both, test.py and wasm-test.py '''

#pylint: disable=too-few-public-methods, c-extension-no-member, broad-except, global-statement, consider-using-with, fixme

from subprocess import check_output, Popen
from time import sleep
from collections import OrderedDict
from sys import stdout
from os.path import exists
from os import remove

import copy
import subprocess
import os
import json
import requests

_OL_DIR = None
_CURR_CONF = None

def setup_config(ol_dir):
    ''' Sets the configuration for the next tests '''

    global _OL_DIR
    _OL_DIR = ol_dir

def get_ol_stats():
    ''' Sets the configuration for the next tests '''

    if os.path.exists(f"{_OL_DIR}/worker/stats.json"):
        with open(f"{_OL_DIR}/worker/stats.json", "r", encoding='utf-8') as statsfile:
            olstats = json.load(statsfile)
        return OrderedDict(sorted(list(olstats.items())))

    return None

def get_worker_output():
    ''' Get the output generated by whatever worker was running '''

    if _OL_DIR is not None:
        with open(os.path.join(_OL_DIR, "worker.out"), "r", encoding='utf-8') as workerfile:
            return workerfile.read().splitlines()
    else:
        # The WebAssembly worker does not create an output currently
        # TODO we should capture it somehow
        return ["_OL_DIR not set"]

def get_current_config():
    ''' Get the current open lambda configuration '''

    return _CURR_CONF

def post(path, data=None):
    ''' Issues a post request to the OL worker '''
    return requests.post(f'http://localhost:5000/{path}',
                         json.dumps(data), timeout=1.0)

def put_conf(conf):
    ''' Sets the specified function and updates the config file on disk '''

    global _CURR_CONF
    with open(os.path.join(_OL_DIR, "config.json"), 'w', encoding='utf-8') as cfile:
        json.dump(conf, cfile, indent=2)
    _CURR_CONF = conf

class TestConf:
    ''' Loads a config and overwrites certain fields with what is set in **keywords '''
    def __init__(self, **keywords):
        self.orig = None

        with open(os.path.join(_OL_DIR, "config.json"), "r", encoding='utf-8') as cfile:
            try:
                self.orig = json.load(cfile)
            except json.JSONDecodeError as err:
                raise ValueError(
                    f"Failed to parse JSON file. Contents are:\n"
                    f"{cfile.read()}"
                ) from err

        new = copy.deepcopy(self.orig)
        for (key, value) in keywords.items():
            if not key in new:
                raise ValueError(f"unknown config param: {key}")

            if isinstance(value, dict):
                for key2 in value:
                    new[key][key2] = value[key2]
            else:
                new[key] = value

        # setup
        put_conf(new)

    def cleanup(self):
        ''' Resets the configuration to what it was before '''
        put_conf(self.orig)

class TestConfContext:
    ''' Keeps track of the OL configuration for a specific test '''

    def __init__(self, **keywords):
        self._conf = None
        self._keywords = keywords

    def __enter__(self):
        self._conf = TestConf(**self._keywords)

    def __exit__(self, _exc_type, _exc_value, _exc_traceback):
        self._conf.cleanup()

def run(cmd):
    ''' Runs a specific command and throws and exception if the command fails '''

    print("RUN", " ".join(cmd))
    try:
        out = check_output(cmd, stderr=subprocess.STDOUT)
        fail = False
    except subprocess.CalledProcessError as err:
        out = err.output
        fail = True

    out = str(out, 'utf-8')

    if not fail:
        if len(out) > 500:
            print(out[:500])
        print(out)
    else:
        print(out)
        raise RuntimeError(f"command ({' '.join(cmd)}) failed")

class DockerWorker():
    ''' Runs OpenLambda with Docker as backend '''

    def __init__(self):
        self._running = False
        self._config = TestConf(sandbox="docker", features={"import_cache": ""})

        try:
            print("Starting Docker container worker")
            run(['sudo', './ol', 'worker', 'up', f'-p={_OL_DIR}', '--detach'])
        except Exception as err:
            raise RuntimeError(f"failed to start worker: {err}") from err

        self._running = True

    def __del__(self):
        self.stop()

    def is_running(self):
        ''' Is the worker (still) running? '''
        return self._running

    @staticmethod
    def name():
        ''' Human-readble name of the worker '''
        return "docker"

    def stop(self):
        ''' Shuts down the worker process '''

        if self.is_running():
            self._running = False
        else:
            return # Already stopped

        try:
            print("Stopping Docker container worker")
            run(['sudo', './ol', 'worker', 'down', '-p='+_OL_DIR])
        except Exception as err:
            raise RuntimeError("Failed to start worker") from err

class SockWorker():
    ''' Runs OpenLambda with SOCK as backend '''

    def __init__(self):
        self._running = False
        self._config = TestConf(sandbox="sock")

        try:
            print("Starting SOCK container worker")
            run(['sudo', './ol', 'worker', 'up', '-p='+_OL_DIR, '--detach'])
        except Exception as err:
            raise RuntimeError(f"failed to start worker: {err}") from err

        self._running = True

    def __del__(self):
        self.stop()

    def is_running(self):
        ''' Is the worker (still) running? '''
        return self._running

    @staticmethod
    def name():
        ''' Human-readble name of the worker '''
        return "SOCK"

    def stop(self):
        ''' Shuts down the worker process '''

        if self.is_running():
            self._running = False
        else:
            return  # Already stopped

        try:
            print("Stopping SOCK container worker")
            run(['sudo', './ol', 'worker', 'down', '-p='+_OL_DIR])
        except Exception as err:
            raise RuntimeError("Failed to start worker") from err


class WasmWorker():
    ''' Runs OpenLambda's WebAssembly worker '''

    def __init__(self):
        # TODO use a similar mechanism to regular OpenLambda

        stdout.write("Starting WebAssembly worker.")
        stdout.flush()

        if exists('./ol-wasm.ready'):
            remove('./ol-wasm.ready')

        self._process = Popen(["./ol-wasm"])

        while not exists('./ol-wasm.ready'):
            sleep(0.5)
            stdout.write('.')
            stdout.flush()

        print("Done")

    def __del__(self):
        self.stop()

    def is_running(self):
        ''' Is the worker (still) running? '''
        return self._process is not None

    @staticmethod
    def name():
        ''' Human-readble name of the worker '''
        return "WebAssembly"

    def stop(self):
        ''' Shuts down the worker process '''

        if not self.is_running():
            return

        print("Stopping WebAssembly worker")
        self._process.terminate()
        self._process = None

def prepare_open_lambda(ol_dir, image="ol-wasm"):
    '''
    Sets up the working director for open lambda,
    and stops currently running worker processes (if any)
    '''
    # init will kill any prior worker and refresh the directory
    # (except for the base "lambda" dir)
    run(['./ol', 'worker', 'init', f'-p={ol_dir}', f'-i={image}'])

def mounts():
    ''' Returns a list of all mounted directories '''

    output = check_output(["mount"])
    output = str(output, "utf-8")
    output = output.split("\n")
    return set(output)


def ol_oom_killer():
    ''' Will terminate OpenLambda if we run out of memory '''

    while True:
        if get_mem_stat_mb('MemAvailable') < 128:
            print("out of memory, trying to kill OL")
            os.system('pkill ol')
        sleep(1)


def get_mem_stat_mb(stat):
    ''' Get the current memory usage in MB '''
    with open('/proc/meminfo', 'r', encoding='utf-8') as memfile:
        for line in memfile:
            if line.startswith(stat+":"):
                parts = line.strip().split()
                assert_eq(parts[-1], 'kB')
                return int(parts[1]) / 1024
    raise ValueError('could not get stat')

def assert_true(val):
    ''' Test helper. Will fail if val is not true '''

    if not val:
        raise ValueError('Expected value to be true')

def assert_eq(actual, expected):
    ''' Test helper. Will fail if actual != expected '''

    if expected != actual:
        raise ValueError(f'Expected value "{expected}", '
                         f'but was "{actual}"')
